import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
import numpy as np

# 创建图形和坐标轴
fig, ax = plt.subplots(figsize=(12, 8))

# 设置背景颜色
ax.set_facecolor('#f0f0f0')
ax.set_xlim(0, 12)
ax.set_ylim(0, 8)

# 绘制矩形框
def draw_rectangle(x, y, width, height, color, label, text_color='black'):
    rect = patches.Rectangle((x, y), width, height, linewidth=1, edgecolor='black', facecolor=color)
    ax.add_patch(rect)
    ax.text(x + width/2, y + height/2, label, ha='center', va='center', fontsize=10, color=text_color)

# 绘制箭头
def draw_arrow(start, end, color='black'):
    ax.annotate('', xy=end, xytext=start, arrowprops=dict(arrowstyle='->', color=color))

# 绘制模块
draw_rectangle(1, 5, 1.5, 1, '#87CEFA', 'MI')
draw_rectangle(1, 1, 1.5, 1, '#87CEFA', 'MS')

draw_rectangle(3, 5, 2, 1, '#FFD700', 'Multi-Head\nAttention')
draw_rectangle(3, 1, 2, 1, '#87CEFA', 'GMP &\nGAP')

draw_rectangle(7, 5, 2, 1, '#90EE90', 'MI')
draw_rectangle(7, 1, 2, 1, '#90EE90', 'MI')

draw_rectangle(9, 4, 1, 1, '#90EE90', 'cat')
draw_rectangle(10.5, 4, 1, 1, '#90EE90', 'FC')

# 绘制箭头
draw_arrow((2.5, 5.5), (3, 5.5))
draw_arrow((2.5, 1.5), (3, 1.5))
draw_arrow((5, 5.5), (7, 5.5))
draw_arrow((5, 1.5), (7, 1.5))

draw_arrow((9, 5.5), (9.5, 5.5))
draw_arrow((9, 1.5), (9.5, 1.5))
draw_arrow((9.5, 5.5), (9.5, 4.5))
draw_arrow((9.5, 1.5), (9.5, 4.5))

draw_arrow((9.5, 4.5), (10, 4.5))
draw_arrow((10, 4.5), (10.5, 4.5))

# 添加文本标签
ax.text(1, 6.5, 'A_img', ha='center', va='center', fontsize=10)
ax.text(1, 0.5, 'A_ske', ha='center', va='center', fontsize=10)
ax.text(10.5, 5, 'weight', ha='center', va='center', fontsize=10)
ax.text(10.5, 3, 'bias', ha='center', va='center', fontsize=10)
ax.text(11.5, 4.5, 'A_f', ha='center', va='center', fontsize=10)

# 设置坐标轴不可见
ax.set_xticks([])
ax.set_yticks([])
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.spines['left'].set_visible(False)

plt.tight_layout()
plt.show()